#include <ros/ros.h>
#include <geometry_msgs/PoseStamped.h>
#include <geometry_msgs/Point.h>
#include <mavros_msgs/CommandBool.h>
#include <mavros_msgs/SetMode.h>
#include <mavros_msgs/State.h>
#include <swarm_plan/FormationInfo.h>
#include <swarm_plan/UavsInfoStamped.h>
#include <Eigen/Eigen>
#include <iostream>

using namespace Eigen;
using namespace std;

double hover_height = 1.5;
bool takeoff = false;

mavros_msgs::State current_state;
Vector3d center_swarm_cur;
Vector2d center_swarm_init;
Vector2d tar_center;
vector<double> control_point;
int max_num = 0;
int count_point = 0;

Matrix<double, 2, 4> swarm_init_pos;
Matrix<double, 2, 4> swarm_cur_pos;

swarm_plan::UavsInfoStamped swarm_info_global;
swarm_plan::FormationInfo formation_tar;

ros::Publisher setpoint_pub[4];
geometry_msgs::PoseStamped pos_tar[4];

// tool func
template<typename T>
void readParam(ros::NodeHandle &nh, std::string param_name, T& loaded_param) {
    const string& node_name = ros::this_node::getName();
    param_name = node_name + "/" + param_name;
    if (!nh.getParam(param_name, loaded_param)) {
        ROS_ERROR_STREAM("Failed to load " << param_name << ", use default value");
    }
    else{
        ROS_INFO_STREAM("Load " << param_name << " success");
    }
}


void loadRosParams(ros::NodeHandle &nh)
{
    readParam<vector<double>>(nh, "control_point", control_point);
    readParam<double>(nh, "hover_height", hover_height);

    max_num = control_point.size() / 2;
    ROS_INFO_STREAM("max_num:" << max_num);

    vector<double> poses(8);
    readParam<vector<double>>(nh, "swarm_init_pos", poses);
    for (int i = 0; i < 4; i++) {
        //对4个无人机依次赋值初始位置
        swarm_init_pos(0, i) = poses[2 * i];
        swarm_init_pos(1, i) = poses[2 * i + 1];
    }
    center_swarm_init(0) = center_swarm_init(0) / 4;
    center_swarm_init(1) = center_swarm_init(1) / 4;
    ROS_INFO_STREAM("center_swarm_init x:" << center_swarm_init(0) << "  y:" << center_swarm_init(1));
}


void state_cb(const mavros_msgs::State::ConstPtr& msg){
    current_state = *msg;
}


void formation_cb(const swarm_plan::FormationInfo::ConstPtr& msg){
    formation_tar = *msg;
    for (int i = 0; i < 4; i++)
    {
        pos_tar[i].pose.position.x = formation_tar.poses[i].x + tar_center(0) - swarm_init_pos(0, i);
        pos_tar[i].pose.position.y = formation_tar.poses[i].y + tar_center(1) - swarm_init_pos(1, i);
        pos_tar[i].pose.position.z = hover_height;
    }
}


void swarm_info_cb(const swarm_plan::UavsInfoStamped::ConstPtr& msg){
    swarm_info_global = *msg;
    center_swarm_cur = Vector3d::Zero();
    for (int i = 0; i < 4; i++)
    {
        center_swarm_cur(0) += swarm_info_global.poses[i].x;
        center_swarm_cur(1) += swarm_info_global.poses[i].y;
        center_swarm_cur(2) += swarm_info_global.poses[i].z;
        swarm_cur_pos(0, i) = swarm_info_global.poses[i].x;
        swarm_cur_pos(1, i) = swarm_info_global.poses[i].y;
    }
    center_swarm_cur(0) = center_swarm_cur(0) / 4;
    center_swarm_cur(1) = center_swarm_cur(1) / 4;
    center_swarm_cur(2) = center_swarm_cur(2) / 4;
    center_swarm_cur(0) += center_swarm_init(0);
    center_swarm_cur(1) += center_swarm_init(1);
}


void pub_pos(){
    for (int i = 0; i < 4;i++){
        pos_tar[i].header.stamp = ros::Time::now();
        setpoint_pub[i].publish(pos_tar[i]);
    }
}

int main (int argc, char** argv) 
{   
    ros::init(argc, argv, "swarm_control_node");

    ros::NodeHandle nh; 
    loadRosParams(nh);

    const int LOOPRATE = 30;
    ros::Rate loop_rate(LOOPRATE);

    ros::Subscriber swarm_info_sub = nh.subscribe("/swarm_plan/swarm_info", 1, swarm_info_cb);
    ros::Subscriber formation_sub = nh.subscribe("/swarm_plan/formation", 1, formation_cb);

    setpoint_pub[0] = nh.advertise<geometry_msgs::PoseStamped>("/uav0/mavros/setpoint_position/local", 1);
    setpoint_pub[1] = nh.advertise<geometry_msgs::PoseStamped>("/uav1/mavros/setpoint_position/local", 1);
    setpoint_pub[2] = nh.advertise<geometry_msgs::PoseStamped>("/uav2/mavros/setpoint_position/local", 1);
    setpoint_pub[3] = nh.advertise<geometry_msgs::PoseStamped>("/uav3/mavros/setpoint_position/local", 1);

    ros::ServiceClient set_mode_client_0 = nh.serviceClient<mavros_msgs::SetMode>("/uav0/mavros/set_mode");
	ros::ServiceClient set_mode_client_1 = nh.serviceClient<mavros_msgs::SetMode>("/uav1/mavros/set_mode");
    ros::ServiceClient set_mode_client_2 = nh.serviceClient<mavros_msgs::SetMode>("/uav2/mavros/set_mode");
    ros::ServiceClient set_mode_client_3 = nh.serviceClient<mavros_msgs::SetMode>("/uav3/mavros/set_mode");

    ros::ServiceClient arming_client_0 = nh.serviceClient<mavros_msgs::CommandBool>("/uav0/mavros/cmd/arming");
    ros::ServiceClient arming_client_1 = nh.serviceClient<mavros_msgs::CommandBool>("/uav1/mavros/cmd/arming");
    ros::ServiceClient arming_client_2 = nh.serviceClient<mavros_msgs::CommandBool>("/uav2/mavros/cmd/arming");
    ros::ServiceClient arming_client_3 = nh.serviceClient<mavros_msgs::CommandBool>("/uav3/mavros/cmd/arming");

    mavros_msgs::CommandBool armed_cmd;
    armed_cmd.request.value = true;
    arming_client_0.call(armed_cmd);
    arming_client_1.call(armed_cmd);
    arming_client_2.call(armed_cmd);
    arming_client_3.call(armed_cmd);
    sleep(2);

    mavros_msgs::SetMode offb_set_mode;
    offb_set_mode.request.custom_mode = "OFFBOARD";
	ros::Time last_request = ros::Time::now();

    for (int i = 0; i < 4;i++){
        pos_tar[i].pose.position.x = 0;
        pos_tar[i].pose.position.y = 0;
        pos_tar[i].pose.position.z = hover_height;
        pos_tar[i].pose.orientation.w = 1.0;
        pos_tar[i].header.stamp = ros::Time::now();
    }

    for(int i = 10; ros::ok() && i > 0; --i){
        pub_pos();
        loop_rate.sleep();
    }
    
    set_mode_client_0.call(offb_set_mode);
    set_mode_client_1.call(offb_set_mode);
    set_mode_client_2.call(offb_set_mode);
    set_mode_client_3.call(offb_set_mode);

	while(ros::ok())
	{
        if(abs(center_swarm_cur(2)-hover_height)<0.4){
            takeoff = true;
        }
        if(!takeoff){
            for (int i = 0; i < 4;i++){
                pos_tar[i].pose.position.x = 0;
                pos_tar[i].pose.position.y = 0;
                pos_tar[i].pose.position.z = hover_height;
                pos_tar[i].pose.orientation.w = 1.0;
                pos_tar[i].header.stamp = ros::Time::now();
            }
            pub_pos();
            ROS_INFO("[SWARM CONTROL]: TAKEOFF !!!!!!!!!!!!!");
        }
        else{
            tar_center(0) = control_point[count_point * 2];
            tar_center(1) = control_point[count_point * 2 + 1];
            if((tar_center-center_swarm_cur.head(2)).norm()<0.30){ 
                count_point++;
                if(count_point>max_num-1){
                    count_point = max_num - 1;
                    ROS_INFO("[SWARM CONTROL]: FINISH!!!!!!!!!!!!!");
                }
            }
            pub_pos();
            ROS_INFO("[SWARM CONTROL]: X: %f Y: %f", tar_center(0), tar_center(1));
        }

        ros::spinOnce();
        loop_rate.sleep();
    }
    return 0;
}
